package org.pathwaycommons.nlp;

import org.cbio.causality.util.FishersExactTest;

import java.util.*;

/**
 * Tests a possible enrichment of gene pairs in the given sets.
 */
public class EnrichmentTester
{
	/**
	 * Sets of gene pairs to test if any of them is enriched.
	 */
	private Map<String, Set<String>> sets;

	/**
	 * Gene pairs to test. Each pair should be represented with a String, where two gene symbols are
	 * separated with a space character.
	 */
	private Set<String> testPairs;

	public EnrichmentTester(Map<String, Set<String>> sets, Set<String> testPairs)
	{
		this.sets = sets;
		this.testPairs = testPairs;
	}

	public Map<String, Double> calcPvals()
	{
		Map<String, Double> pvals = new HashMap<String, Double>();

		int backgroundSize = getBackgroundPairs().size();
		int selectedSize = testPairs.size();

		for (String feature : sets.keySet())
		{
			int featuredSize = sets.get(feature).size();
			int featuredSelectedSize = overlapSize(sets.get(feature), testPairs);

			pvals.put(feature, FishersExactTest.calcEnrichmentPval(
				backgroundSize, featuredSize, selectedSize, featuredSelectedSize));
		}

		return pvals;
	}

	private Set<String> getBackgroundGenes()
	{
		Set<String> genes = new HashSet<String>();
		for (Set<String> pairs : sets.values())
		{
			for (String pair : pairs)
			{
				Collections.addAll(genes, pair.split(" "));
			}
		}
		for (String testPair : testPairs)
		{
			Collections.addAll(genes, testPair.split(" "));
		}
		return genes;
	}

	private Set<String> getBackgroundPairs()
	{
		Set<String> pairs = new HashSet<String>();
		Set<String> genes = getBackgroundGenes();
		for (String g1 : genes)
		{
			for (String g2 : genes)
			{
				if (g1.equals(g2)) continue;

				pairs.add(g1 + " " + g2);
			}
		}
		return pairs;
	}

	private int overlapSize(Set<String> s1, Set<String> s2)
	{
		Set<String> set = new HashSet<String>(s2);
		set.retainAll(s1);
		return set.size();
	}

}
